import os
import sys
import argparse
import numpy as np
import mindspore

from mindspore import Tensor,export,load_checkpoint,context
from mind3d.utils.load_yaml import load_yaml
from mind3d.models.dgcnn import DGCNN_cls
from mind3d.models.pointnet import PointNet_cls
from mind3d.models.pointnet2 import Pointnet2clsModelSSG
from mind3d.models.PointTransformer import PointTransformerCls
from mindspore import context, load_checkpoint, load_param_into_net

if not os.path.exists("./mindir"):
    os.mkdir("./mindir")

num_classes=40

def export(opt,args):
    if opt.model_name=="DGCNN_cls":
        classifier=DGCNN_cls(opt,output_channels=opt['train'].get("num_classes"))
        load_checkpoint(args.model,net=classifier)
        input_data=Tensor(np.random.uniform(0.0,1.0,size=[32,1024,3]),mindspore.float32)
        export(classifier,input_data,file_name="./mindir/dgcnn_cls",file_format=args.file_format)
        print("successfully export model")
    elif opt.model_name == "PointNet_cls":
        classifier = PointNet_cls(k=num_classes)
        load_checkpoint(args.model, net=classifier)
        input_data = Tensor(np.random.uniform(0.0, 1.0, size=[32, 1024, 3]), mindspore.float32)
        export(classifier, input_data, file_name='./mindir/pointnet_cls', file_format=args.file_format)
        print("successfully export model")
    elif opt.model_name == "Pointnet2clsModelSSG":
        classifier = Pointnet2clsModelSSG(normal_channel=False)
        load_checkpoint(args.model, net=classifier)
        input_data = Tensor(np.random.uniform(0.0, 1.0, size=[16, 1024, 3]), mindspore.float32)
        export(classifier, input_data, file_name='./mindir/pointnet2_cls', file_format=args.file_format)
        print("successfully export model")
    elif opt.model_name=="PointTransformerCls":
        network = PointTransformerCls()
        pretrain_ckpt_path = opt['Export'].get("ckpt_file")
        file_name = "PointTransformerCls"
        file_format = "MINDIR"
        input_shape = [opt['Export'].get("batch_size"), opt['Export'].get("num_points"), 6]
        load_param_into_net(network, load_checkpoint(pretrain_ckpt_path))
        input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shape).astype(np.float32))
        export(network, input_array, file_name=file_name, file_format=file_format)
        print(f"Successful export model.")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Mindspore DGCNN")
    parser.add_argument("--device_target", default="GPU", help='device')
    parser.add_argument('--model', type=str, default="/home/cxh/音乐/MindSpore-DGCNN/pretrain/seg/model_6.ckpt", help="model path")
    parser.add_argument("--file_format", type=str, default="MINDIR", help="export file format")
    parser.add_argument("-opt", type=str, default="/home/cxh/音乐/MindSpore-DGCNN/dgcnn_s3dis_seg.yaml",
                        help='Path to option YAML file.')
    args = parser.parse_known_args()[0]
    os.environ["CUDA_VISBLE_DEVICES"] = '0'
    context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target)
    opt = load_yaml(args.opt)
    export(opt,args)


